from exper.Experiment import Experiment
import torch


class ExpMUAPs(Experiment):

    def __init__(self, config, dataer, neter):

        super(ExpMUAPs, self).__init__('MUAPs', config)
        self.dataer = dataer
        self.neter = neter

    def attck_acc(self, atk, seed, repeat_times=3, default_num=500):

        Attack_numbers_list = [0 for i in range(self.config['base_size'])]
        
        loader = self.dataer.get_shuffle_loader(seed=seed, isTrain=False, batch_size=1)
        
        count = 0
        succ_number = 0
        for (image, label) in loader:
            
            if count == default_num:
                break
            
            test_image = image.to(self.neter.device)
            test_label = label.to(self.neter.device)
            output = self.neter.net(test_image)
            _, pre = torch.max(output.data, 1)
            
            if pre != test_label:
                print('Pass the error sample')
                continue
            
            succ_list = atk.untarget_UAP_show(
                            origin_sample=image,
                            source_class=label[0],
                            epsilon=self.config['epsilon'],
                            norm_type=self.config['attck_norm'],
                            method=self.config['attck_method'],
                            batch=self.config['base_batch'],
                            size=self.config['base_size'],
                            mode=self.config['base_mode'],
                        )
            count += 1
            
            ## exclude index 9
            new_succ_list = [item for item in succ_list if item != 9]
            succ_list = new_succ_list

            if len(succ_list) > 0:
                succ_number += 1
            
            for index in succ_list:
                Attack_numbers_list[index] += 1
        print(succ_number)
        print(Attack_numbers_list)
